from openai import OpenAI

OPENAI_API_KEY = "d52e49a1-36ea-44bb-bc6e-65ce789a72f6"
OPENAI_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
OPENAI_MODEL = "doubao-seed-1-6-250615"


# 非流式输出
def doubao_qa(prompt):
    client = OpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)
    try:
        resp = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
        )
        return resp.choices[0].message.content
    except Exception as e:
        return f"[Doubao接口调用失败] {e}"


# 流式输出（生成器）
def doubao_qa_stream(prompt):
    client = OpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)
    try:
        stream = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
            stream=True,
        )
        for chunk in stream:
            yield chunk.choices[0].delta.content
    except Exception as e:
        yield f"[Doubao接口调用失败] {e}"


# 示例用法
if __name__ == "__main__":
    question = "鲁迅和周树人什么关系?"
    print("--- 非流式 ---")
    # answer = doubao_qa(question)
    # print("Doubao回答：", answer)
    print("--- 流式 ---")
    for part in doubao_qa_stream(question):
        print(part, end="", flush=True)
    print()
